import enum
import mindspore as ms
import mindspore.nn as nn
import random
# nn.utils 以及 nn.functional在ms中是分散的，具体用到的时候再修改成对应的值
from collections import OrderedDict

from setuptools import SetuptoolsDeprecationWarning
from .core import Model, LayerNorm
from .modules import *
from .spectral_norm import SpectralNorm

class Discriminator(Model):

    def __init__(self, ninputs, fmaps,
                 kwidth, poolings,
                 pool_type='none',
                 pool_slen=None,
                 norm_type='bnorm',
                 bias=True,
                 phase_shift=None, 
                 sinc_conv=False):
        super().__init__(name='Discriminator')
        # phase_shift randomly occurs within D layers
        # as proposed in https://arxiv.org/pdf/1802.04208.pdf
        # phase shift has to be specified as an integer
        self.phase_shift = phase_shift
        if phase_shift is not None:
            assert isinstance(phase_shift, int), type(phase_shift)
            assert phase_shift > 1, phase_shift
        if pool_slen is None:
            raise ValueError('Please specify D network pool seq len '
                             '(pool_slen) in the end of the conv '
                             'stack: [inp_len // (total_pooling_factor)]')
        ninp = ninputs
        # SincNet as proposed in 
        # https://arxiv.org/abs/1808.00158
        if sinc_conv:
            # build sincnet module as first layer
            self.sinc_conv = SincConv(fmaps[0] // 2,
                                      251, 16e3, padding='SAME')
            inp = fmaps[0]
            fmaps = fmaps[1:]
        self.enc_blocks = nn.CellList()
        for pi, (fmap, pool) in enumerate(zip(fmaps, poolings), start=1):
            enc_block = GConv1DBlock(ninp, fmap, kwidth, stride=pool, bias=bias, norm_type=norm_type)
            self.enc_blocks.append(enc_block)
            ninp = fmap
        self.pool_type = pool_type
        if pool_type == 'none':
            # resize tensor to fit into FC directly
            pool_slen *= fmaps[-1]
            self.fc = nn.SequentialCell(
                nn.Dense(pool_slen, 256),
                nn.PReLU(256),
                nn.Dense(256, 128),
                nn.PReLU(128),
                nn.Dense(128, 1)
            )
            if norm_type == 'snorm':
                self.fc[0] = SpectralNorm(self.fc[0])
                self.fc[2] = SpectralNorm(self.fc[2])
                self.fc[3] = SpectralNorm(self.fc[3])
        elif pool_type == 'conv':
            self.pool_conv = nn.Conv1d(fmap[-1], 1, 1, pad_mode="pad")
            self.fc = nn.Dense(pool_slen, 1)
            if norm_type == 'snorm':
                self.pool_conv = SpectralNorm(self.pool_conv)
                self.fc = SpectralNorm(self.fc)
        elif pool_type == 'gmax':
            self.gmax = nn.AdaptiveMaxPool1d(1)
            self.fc = nn.Dense(fmap[-1], 1, 1)
            if norm_type == 'snorm':
                self.fc = SpectralNorm(self.fc)
        elif pool_type == 'gavg':
            self.gavg = nn.AdaptiveAvgPool1d(1)
            self.fc = nn.Dense(fmaps[-1], 1, 1)
            if norm_type == 'snorm':
                self.fc = SpectralNorm(self.fc)
        elif pool_type == 'mlp':
            self.mlp = nn.SequentialCell(
                nn.Conv1d(fmaps[-1], fmap[-1], 1, pad_mode='pad'),
                nn.PReLU(fmaps[-1]),
                nn.Conv1d(fmaps[-1], 1, 1, pad_mode='pad')
            )
            if norm_type == 'snorm':
                self.mlp[0] = SpectralNorm(self.mlp[0])
                self.mlp[1] = SpectralNorm(self.mlp[1])
        else:
            raise TypeError('Unrecognized pool type: ', pool_type)
    
    def construct(self, x):
        h = x
        if hasattr(self, 'sinc_conv'):
            h_l, h_r = ms.ops.split(h, 1, 2)
            h_l = self.sinc_conv(h_l)
            h_r = self.sinc_conv(h_r)
            h = ms.ops.concat((h_l, h_r), axis=1)
            # store intermediate activations
        # int_act = {}
        for ii, layer in enumerate(self.enc_blocks):
            # if self.phase_shift is not None:
            #     shift = random.randint(1, self.phase_shift)
            #     # 0.5 chance of shifting right or left
            #     right = random.random() > 0.5
            #     # split tensor in time dim (dim 2)
            #     if right:
            #         sp1 = h[:, :, :h.shape[2] - shift]
            #         sp2 = h[:, :, h.shape[2] - shift:]
            #         h = ms.ops.concat((sp2, sp1), dim=2)
            #     else:
            #         sp1 = h[:, :, :shift]
            #         sp2 = h[:, :, shift:]
            #         h = ms.ops.concat((sp2, sp1), dim=2)
            h = layer(h)
            # int_act['h_{}'.format(ii)] = h
        y = None
        if self.pool_type == 'conv':
            h = self.pool_conv(h)
            h = h.view(h.size(0), -1)
            # int_act['avg_conv_h'] = h
            y = self.fc(h)
        elif self.pool_type == 'none':
            h = h.view(h.size(0), -1)
            y = self.fc(h)
        elif self.pool_type == 'gmax':
            h = self.gmax(h)
            h = h.view(h.size(0), -1)
            y = self.fc(h)
        elif self.pool_type == 'gavg':
            h = self.gavg(h)
            h = h.view(h.size(0), -1)
            y = self.fc(h)
        elif self.pool_type == 'mlp':
            y = self.mlp(h)
        # int_act['logit'] = y
        return y


if __name__ == '__main__':
    # pool_slen = 16 because we have input len 16384
    # and we perform 5 pooling layers of 4, so 16384 // (4 ** 5) = 16
    disc = Discriminator(2, [64, 128, 256, 512, 1024],
                         31, [4] * 5, pool_type='none',
                         pool_slen=16)
    print(disc)
    print('Num params: ', disc.get_n_params())
    x = ms.numpy.randn(1, 2, 16384)
    y, _ = disc(x)
    print(y)
    print('x size: {} -> y size: {}'.format(x.size(), y.size()))


